# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####


bl_info = {
  "name": "Export a custom skeleton format",
  "description": "work in progress",
  "author": "Thibault Coppex",
  "version": (0, 2),
  "blender": (2, 64, 0),
  "location": "File > Export > Skeleton SKL format",
  "category": "Import-Export"
}


"""
  NOTES
    ..
  
  TODO
    ..
  
  CREDITS
    io_export_unreal_psk_psa.py
"""


import os
import bpy
import operator
import mathutils
from bpy.props import *



# ================================================= #
# CONSTANTS
# ================================================= #

FILE_EXT = ".skl"


# ================================================= #
# STRUCTURES
# ================================================= #


class Bone():
  
  def __init__(self, name="", id=0, qRot=mathutils.Quaternion(), vPos=mathutils.Vector()):
    self.Name = name[:32]
    self.ParentId = id
    self.Rotation = qRot
    self.Position = vPos


  def __str__(self):
      return "%s %d %.3f %.3f %.3f %.3f %.3f %.3f %.3f" %  \
            (self.Name, self.ParentId, \
             self.Rotation.w, self.Rotation.x, self.Rotation.y, self.Rotation.z,\
             self.Position.x, self.Position.y, self.Position.z )

      

class SKLData():
  
  def __init__(self):
    self.Bones = []
    
  def addBone(self, bone):
    self.Bones.append( bone )
    return len(self.Bones)-1

  
  

# ================================================= #
# UTILS
# ================================================= #

def get_armature():
  
  bpy.ops.object.mode_set(mode='OBJECT')
    
  context = bpy.context
  
  # Return directly the armature if selected
  if context.active_object != None and context.active_object.type == 'ARMATURE':
    return context.active_object
  
  # Or search for a matching one
  arm = [ obj for obj in context.scene.objects if obj.type == 'ARMATURE' ]
  
  if len(arm) != 1:
    raise Exception( "Error : a single armature must be selected" )
  
  return arm[0]



def export_armature(armature, sklData):  
  # Search for root candidates
  root_candidates = [b for b in armature.data.bones if b.parent == None and b.name[:4] == "Root"]
  
  # Only one node can be eligible
  if len(root_candidates) != 1:
    raise Exception( "A single root must exist in the armature." )
  
  # Get the root
  root = root_candidates[0]
  del root_candidates  
  
  
  #Matrix to transform the skeleton in OpenGL space coordinate
  # (Blender X, Y, Z is OpenGL X, Z, -Y so we rotate -90deg around x axis)
  qBlenderToGL = mathutils.Quaternion( (0.707, -0.707, 0.0, 0.0) )
  
  # Root orientation & position
  rotMatrix = root.matrix * armature.matrix_local.to_3x3()
  qRotation = rotMatrix.to_quaternion()
  qRotation *= qBlenderToGL
  vPosition = armature.matrix_local * root.head  
  
  #return the root id, which is 0
  root_id = sklData.addBone( Bone( root.name, -1, qRotation, vPosition) )
    
  # recursively compute others bones data
  for child in root.children:
    recursively_add_bone( child, root_id, sklData)
  

def recursively_add_bone( bone, parent_id, sklData):   
  qRotation = bone.matrix.to_quaternion()
  
  qRotParent = bone.parent.matrix.to_quaternion().inverted()
  offsetParent = qRotParent * (bone.parent.tail - bone.parent.head)      
  vPosition = offsetParent + bone.head
  
  bone_id = sklData.addBone( Bone( bone.name, parent_id, qRotation, vPosition) )
  
  for child in bone.children:
    recursively_add_bone( child, bone_id, sklData)



# ================================================= #
# EXPORTER
# ================================================= #

def Save(filepath=""):
  
  sklData = SKLData()
  
  armature = get_armature()  
  export_armature( armature, sklData )  

  file = open( filepath, "w" );
  
  file.write( str(len(sklData.Bones)) + '\n' )
  
  # JointName JointParentID Qw Qx Qy Qz Tx Ty Tz
  for bone in sklData.Bones:
    file.write( str(bone)+'\n' )

  file.close()
  
  

class ExportSKL(bpy.types.Operator):
  __doc__ = ""
  
  bl_idname = "export_skeleton.skl"
  bl_label = 'Export SKL'
  
  filepath = StringProperty(name="File Path", maxlen=1024, default="")
  
  def execute( self, context):
    Save( self.filepath )
    return {'FINISHED'}
  
  


# ================================================= #
# REGISTER / UNREGISTER
# ================================================= #

def menu_func(self, context):
  default_path = os.path.splitext( bpy.data.filepath )[0] + FILE_EXT
  self.layout.operator( ExportSKL.bl_idname, text="Skeleton data ("+FILE_EXT+")").filepath = default_path

def register():
  bpy.utils.register_module(__name__)
  bpy.types.INFO_MT_file_export.append(menu_func)

def unregister():
  bpy.utils.unregister_module(__name__)
  bpy.types.INFO_MT_file_export.remove(menu_func)


if __name__ == "__main__":
  Save( "/data/Code/projets/aer-anima/data/simple.skl" )
  #register()
  #unregister()

